import os
import glob
import copy
import cv2
import pandas as pd
import numpy as np
import seaborn as sns
import random
from PIL import Image
import matplotlib.pyplot as plt

try:
    import xml.etree.CElementTree as ET
except:
    import xml.etree.ElementTree as ET

def mkPath(paths):
    if not isinstance(paths,list):
        paths = list(paths)
    for path in paths:
        if os.path.isfile(path):
            path = os.path.dirname(path)
        if not os.path.exists(path):
            os.makedirs(path)

def copy_bbox(xmlFilePath,bbox):
    image_path = 'data_enhance/image'
    xml_path = 'data_enhance/xml'
    mkPath([image_path,xml_path])
    
    module_input_size = 768
    # 计算copybbox的中心偏移,并复制图片
    fileName = os.path.splitext(os.path.basename(xmlFilePath))[0]
    img = Image.open(os.path.join('JPEGImages',fileName+'.jpg')) #w,h
    img_w,img_h = img.size
    ratio_xw = img_w*1.0/768
    ratio_yh = img_h*1.0/768
    bbox_w,bbox_h,bbox_x,bbox_y = int(bbox[0]*ratio_xw),int(bbox[1]*ratio_yh),int(bbox[2]*ratio_xw),int(bbox[3]*ratio_yh)
    x1,y1,x2,y2 = int(bbox_x - bbox_w//2),int(bbox_y - bbox_h//2),int(bbox_x + bbox_w//2),int(bbox_y + bbox_h//2)
    #print(x1,y1,x2,y2)
    # bbox 复制
    copy_bbox_l_x1,copy_bbox_l_x2,copy_bbox_r_x1,copy_bbox_r_x2  = 0,0,0,0
    if bbox_x-bbox_w > bbox_w:
        copy_bbox_l_x = np.random.randint(bbox_w,bbox_x-bbox_w,1 )[0]
        #nonlocal copy_bbox_l_x1,copy_bbox_l_x2 #使用外部闭包中的变量（函数嵌套时的外部函数变量）
        copy_bbox_l_x1,copy_bbox_l_x2 = copy_bbox_l_x - bbox_w//2,copy_bbox_l_x + bbox_w//2
    if bbox_x+ bbox_w < img_w - bbox_w:
        copy_bbox_r_x = np.random.randint(bbox_x+ bbox_w, img_w - bbox_w,1 )[0]
        #nonlocal copy_bbox_r_x1,copy_bbox_r_x2
        copy_bbox_r_x1,copy_bbox_r_x2 = copy_bbox_r_x - bbox_w//2,copy_bbox_r_x + bbox_w//2
        #PIl 的bbox坐标为 (x1,y1)或(x1,y1,x2,y2)的形式,当给定四元组时，图片大小必须匹配
    enhance_img = img
    crop = img.crop((x1,y1,x2,y2))
    #crop.show()
    if copy_bbox_l_x1 > 0:
        enhance_img.paste(crop,(copy_bbox_l_x1,y1))
    if copy_bbox_r_x1 > 0:
        enhance_img.paste(crop,(copy_bbox_r_x1,y1))
    enhance_img.convert('RGB').save(f'{image_path}/{fileName}_enhance.jpg')
    #enhance_img.show()
    
    # 保存复制bbox到xml
    tree = ET.parse(xmlFilePath)
    root = tree.getroot()
    if copy_bbox_l_x1 > 0:
        copy_l = root.find('object')
        copy_l = copy.deepcopy(copy_l)
        copy_l.find('bndbox/xmin').text = str(copy_bbox_l_x1)
        copy_l.find('bndbox/ymin').text = str(y1)
        copy_l.find('bndbox/xmax').text = str(copy_bbox_l_x2)
        copy_l.find('bndbox/ymax').text = str(y2)
        root.append(copy_l)
        tree.write(f'{xml_path}/{fileName}_enhance.xml')
    if copy_bbox_r_x1 > 0:
        copy_r = root.find('object')
        copy_r = copy.deepcopy(copy_r)
        copy_r.find('bndbox/xmin').text = str(copy_bbox_r_x1)
        copy_r.find('bndbox/ymin').text = str(y1)
        copy_r.find('bndbox/xmax').text = str(copy_bbox_r_x2)
        copy_r.find('bndbox/ymax').text = str(y2)
        root.append(copy_r)
    tree.write(f'{xml_path}/{fileName}_enhance.xml')
    
    #test = ET.parse(f'{xml_path}/{fileName}_enhance.xml')
    #print(ET.tostring(test.getroot()))


def bbox_from_xml():
    module_input_size = 768
    wh_list = []
    center_list = []
    for xmlFilePath in glob.glob("Annotations/*.xml"):
        tree = ET.parse(xmlFilePath)
        root = tree.getroot()
        ImageSize_w = int(root.find('size').find('width').text)
        ImageSize_h = int(root.find('size').find('height').text)
        if not ImageSize_w or not ImageSize_h:
            filename = root.find('filename').text
            image = cv2.imread(f"JPEGImages/{filename}")
            ImageSize_w = image.shape[1]
            ImageSize_h = image.shape[0]
        for bbox in root.iter('bndbox'):
            xmax = int(bbox.find('xmax').text)
            xmin = int(bbox.find('xmin').text)
            ymax = int(bbox.find('ymax').text)
            ymin = int(bbox.find('ymin').text)
            w = (xmax - xmin)*module_input_size/ImageSize_w
            h = (ymax - ymin)*module_input_size/ImageSize_h
            center_x = (xmax + xmin)*1.0/2*module_input_size/ImageSize_w
            center_y = (ymax + ymin)*1.0/2*module_input_size/ImageSize_h
            if w and h:
                wh_list.append([w,h])
                center_list.append([center_x,center_y])
                        #数据集增强
            if center_y < module_input_size -600 and pow(w*h,0.25) < 10:
                #print(xmin,ymin,xmax,ymax)
                copy_bbox(xmlFilePath,[w,h,center_x,center_y])
                
    wh_list,center_list = np.array(wh_list),np.array(center_list)
    bbox = np.concatenate((wh_list,center_list),axis=1)
    with open('bbox.txt','a+') as f:
        np.savetxt(f,bbox,fmt='%d,%d,%d,%d',delimiter=',')
    
    return bbox #w,h,cente_x,cente_y

bbox_from_xml()  